package turritopsis

import (
    "context"
    "log"
    "math/rand"
    "sync"
    "sync/atomic"
    "time"
    "adkr/pkg/adkr"
)

// RealADKRConfig represents configuration for real ADKR
type RealADKRConfig struct {
	N           int           // Number of nodes
	F           int           // Number of faulty nodes
    ACSSDelay   time.Duration // ACSS processing delay per instance
	MVBADelay   time.Duration // MVBA processing delay
	KeyGenDelay time.Duration // Key generation delay
    // Parallelization & control
    ACSSParallel int           // Max parallel ACSS instances (0 or <1 means N)
    ACSSTimeout  time.Duration // Optional timeout for ACSS phase (0 means no timeout)
    MVBATimeout  time.Duration // Optional timeout for MVBA phase
    KeyGenTimeout time.Duration // Optional timeout for KeyGen phase
    // Crypto simulation params
    Curve       string // "secp256k1" or "bls12-381"
    VerifyBatch bool
    VerifyCount int // number of signature verifications to simulate
}

// RealADKR implements a realistic ADKR simulation
type RealADKR struct {
	config RealADKRConfig
	logger *log.Logger
	thresholdCrypto *ThresholdCrypto
	metrics *Metrics
}

// NewRealADKR creates a new real ADKR instance
func NewRealADKR(config RealADKRConfig, logger *log.Logger) *RealADKR {
	if logger == nil {
		logger = log.Default()
	}
	
	// Initialize threshold crypto
	thresholdConfig := ThresholdCryptoConfig{
		Curve:       config.Curve,
		VerifyBatch: config.VerifyBatch,
		VerifyCount: config.VerifyCount,
	}
	thresholdCrypto := NewThresholdCrypto(thresholdConfig)
	
	return &RealADKR{
		config: config, 
		logger: logger,
		thresholdCrypto: thresholdCrypto,
		metrics: &Metrics{},
	}
}

// SetMetrics sets the metrics instance for detailed recording
func (r *RealADKR) SetMetrics(metrics *Metrics) {
	r.metrics = metrics
}

// RunADKR executes the real ADKR protocol
func (r *RealADKR) RunADKR(epoch int) (*adkr.KeyRefreshOutput, error) {
	r.logger.Printf("Real ADKR[epoch-%d]: Starting with N=%d, F=%d", epoch, r.config.N, r.config.F)
	
	// Phase 1: ACSS - Asynchronous Complete Secret Sharing
	// Run n parallel ACSS instances
	r.logger.Printf("Real ADKR[epoch-%d]: Starting %d ACSS instances", epoch, r.config.N)
    acssStart := time.Now()
    
    // Record ACSS phase start
    if r.metrics != nil {
        r.metrics.RecordEpochPhase(epoch, "acss", acssStart, time.Time{})
    }

    // Context with optional timeout for ACSS
    var (
        ctx    context.Context
        cancel context.CancelFunc
    )
    if r.config.ACSSTimeout > 0 {
        ctx, cancel = context.WithTimeout(context.Background(), r.config.ACSSTimeout)
    } else {
        ctx, cancel = context.WithCancel(context.Background())
    }
    defer cancel()

    // Bounded parallelism via semaphore
    parallel := r.config.ACSSParallel
    if parallel <= 0 || parallel > r.config.N {
        parallel = r.config.N
    }
    sem := make(chan struct{}, parallel)
    var wg sync.WaitGroup
    var firstErr atomic.Value // store error

    perInst := func(i int) {
        defer wg.Done()
        // acquire permit (backpressure)
        select {
        case sem <- struct{}{}:
            defer func() { <-sem }()
        case <-ctx.Done():
            return
        }

        // simulate per-instance processing with jitter
        jitter := time.Duration(rand.Intn(int(r.config.ACSSDelay/2 + 1)))
        delay := r.config.ACSSDelay + jitter

        select {
        case <-time.After(delay):
            // ok
        case <-ctx.Done():
            return
        }
    }

    wg.Add(r.config.N)
    for i := 0; i < r.config.N; i++ {
        go perInst(i)
    }
    done := make(chan struct{})
    go func() { wg.Wait(); close(done) }()

    select {
    case <-done:
        // all ACSS instances finished
    case <-ctx.Done():
        if ctx.Err() != nil && firstErr.Load() == nil {
            firstErr.Store(ctx.Err())
        }
    }

    acssEnd := time.Now()
    acssDuration := acssEnd.Sub(acssStart)
	r.logger.Printf("Real ADKR[epoch-%d]: ACSS completed in %v", epoch, acssDuration)
	
	// Record ACSS phase end and message counts
	if r.metrics != nil {
		r.metrics.RecordEpochPhase(epoch, "acss", acssStart, acssEnd)
		// Simulate message counts for ACSS phase
		acssMessages := r.config.N * 3 // Each ACSS instance sends ~3 messages
		messageCounts := MessageCounts{
			ACSSMessages: acssMessages,
			TotalMessages: acssMessages,
			BytesSent: int64(acssMessages * 256), // ~256 bytes per message
			BytesReceived: int64(acssMessages * 256),
		}
		r.metrics.RecordEpochMessages(epoch, messageCounts)
	}
	
    // Phase 2: MVBA - Multi-Validated Byzantine Agreement (with optional timeout)
    r.logger.Printf("Real ADKR[epoch-%d]: Starting MVBA", epoch)
    mvbaStart := time.Now()
    
    // Record MVBA phase start
    if r.metrics != nil {
        r.metrics.RecordEpochPhase(epoch, "mvba", mvbaStart, time.Time{})
    }
    if r.config.MVBATimeout > 0 {
        mvbaCtx, mvbaCancel := context.WithTimeout(context.Background(), r.config.MVBATimeout)
        defer mvbaCancel()
        mvbaDelay := r.config.MVBADelay + time.Duration(rand.Intn(int(r.config.MVBADelay/2+1)))
        select {
        case <-time.After(mvbaDelay):
        case <-mvbaCtx.Done():
            r.logger.Printf("Real ADKR[epoch-%d]: MVBA timed out after %v", epoch, r.config.MVBATimeout)
        }
    } else {
        mvbaDelay := r.config.MVBADelay + time.Duration(rand.Intn(int(r.config.MVBADelay/2+1)))
        time.Sleep(mvbaDelay)
    }
    mvbaEnd := time.Now()
    mvbaDuration := mvbaEnd.Sub(mvbaStart)
	r.logger.Printf("Real ADKR[epoch-%d]: MVBA completed in %v", epoch, mvbaDuration)
	
	// Record MVBA phase end and update message counts
	if r.metrics != nil {
		r.metrics.RecordEpochPhase(epoch, "mvba", mvbaStart, mvbaEnd)
		// Update message counts for MVBA phase
		mvbaMessages := r.config.N * 2 // MVBA typically requires 2 rounds
		// Get existing message counts and update
		existingCounts := MessageCounts{}
		if len(r.metrics.EpochsExt) > epoch {
			existingCounts = r.metrics.EpochsExt[epoch].MessageCounts
		}
		existingCounts.MVAMessages = mvbaMessages
		existingCounts.TotalMessages += mvbaMessages
		existingCounts.BytesSent += int64(mvbaMessages * 512) // ~512 bytes per MVBA message
		existingCounts.BytesReceived += int64(mvbaMessages * 512)
		r.metrics.RecordEpochMessages(epoch, existingCounts)
	}
	
    // Phase 3: Key Generation - Threshold signature generation (with optional timeout)
    r.logger.Printf("Real ADKR[epoch-%d]: Starting key generation", epoch)
    keyGenStart := time.Now()
    
    // Record KeyGen phase start
    if r.metrics != nil {
        r.metrics.RecordEpochPhase(epoch, "keygen", keyGenStart, time.Time{})
    }
    
    // Generate threshold keys
    thresholdPK, publicKeys, _, err := r.thresholdCrypto.GenerateThresholdKeys(r.config.N, r.config.F+1)
    if err != nil {
        r.logger.Printf("Real ADKR[epoch-%d]: Key generation failed: %v", epoch, err)
        return nil, err
    }
    
    if r.config.KeyGenTimeout > 0 {
        kgCtx, kgCancel := context.WithTimeout(context.Background(), r.config.KeyGenTimeout)
        defer kgCancel()
        keyGenDelay := r.config.KeyGenDelay + time.Duration(rand.Intn(int(r.config.KeyGenDelay/2+1)))
        select {
        case <-time.After(keyGenDelay):
        case <-kgCtx.Done():
            r.logger.Printf("Real ADKR[epoch-%d]: KeyGen timed out after %v", epoch, r.config.KeyGenTimeout)
        }
    } else {
        keyGenDelay := r.config.KeyGenDelay + time.Duration(rand.Intn(int(r.config.KeyGenDelay/2+1)))
        time.Sleep(keyGenDelay)
    }
    keyGenEnd := time.Now()
    keyGenDuration := keyGenEnd.Sub(keyGenStart)
	r.logger.Printf("Real ADKR[epoch-%d]: Key generation completed in %v (thresholdPK=%v, keys=%d)", epoch, keyGenDuration, thresholdPK != nil, len(publicKeys))
	
	// Record KeyGen phase end and update message counts
	if r.metrics != nil {
		r.metrics.RecordEpochPhase(epoch, "keygen", keyGenStart, keyGenEnd)
		// Update message counts for KeyGen phase
		keyGenMessages := r.config.N // Each node generates and shares keys
		// Get existing message counts and update
		existingCounts := MessageCounts{}
		if len(r.metrics.EpochsExt) > epoch {
			existingCounts = r.metrics.EpochsExt[epoch].MessageCounts
		}
		existingCounts.KeyGenMessages = keyGenMessages
		existingCounts.TotalMessages += keyGenMessages
		existingCounts.BytesSent += int64(keyGenMessages * 128) // ~128 bytes per key share
		existingCounts.BytesReceived += int64(keyGenMessages * 128)
		r.metrics.RecordEpochMessages(epoch, existingCounts)
	}

    // Phase 4: Threshold signature verifications (real batch verification)
    if r.config.VerifyCount > 0 {
        r.logger.Printf("Real ADKR[epoch-%d]: Starting batch verification of %d signatures", epoch, r.config.VerifyCount)
        verifyStart := time.Now()
        
        // Record Verify phase start
        if r.metrics != nil {
            r.metrics.RecordEpochPhase(epoch, "verify", verifyStart, time.Time{})
        }
        
        // Create signature batch for testing
        batch, err := r.thresholdCrypto.CreateSignatureBatch(r.config.VerifyCount)
        if err != nil {
            r.logger.Printf("Real ADKR[epoch-%d]: Failed to create signature batch: %v", epoch, err)
        } else {
            // Perform batch verification
            valid, duration, err := r.thresholdCrypto.BatchVerifySignatures(batch)
            if err != nil {
                r.logger.Printf("Real ADKR[epoch-%d]: Batch verification failed: %v", epoch, err)
            } else {
                r.logger.Printf("Real ADKR[epoch-%d]: Batch verification completed: valid=%v, duration=%v (curve=%s, batch=%v)", 
                    epoch, valid, duration, r.config.Curve, r.config.VerifyBatch)
            }
        }
        
        verifyEnd := time.Now()
        verifyDuration := verifyEnd.Sub(verifyStart)
        r.logger.Printf("Real ADKR[epoch-%d]: Total verification phase completed in %v", epoch, verifyDuration)
        
        // Record Verify phase end and update message counts
        if r.metrics != nil {
            r.metrics.RecordEpochPhase(epoch, "verify", verifyStart, verifyEnd)
            // Update message counts for Verify phase
            verifyMessages := r.config.VerifyCount / 10 // Batch verification reduces message count
            // Get existing message counts and update
            existingCounts := MessageCounts{}
            if len(r.metrics.EpochsExt) > epoch {
                existingCounts = r.metrics.EpochsExt[epoch].MessageCounts
            }
            existingCounts.VerifyMessages = verifyMessages
            existingCounts.TotalMessages += verifyMessages
            existingCounts.BytesSent += int64(verifyMessages * 64) // ~64 bytes per verification result
            existingCounts.BytesReceived += int64(verifyMessages * 64)
            r.metrics.RecordEpochMessages(epoch, existingCounts)
        }
    }
	
	totalDuration := time.Since(acssStart)
	r.logger.Printf("Real ADKR[epoch-%d]: Total execution time: %v", epoch, totalDuration)
	
	// Return realistic KeyRefreshOutput
	return &adkr.KeyRefreshOutput{
		ConfigChain: [][]int{{0, 1, 2, 3}}, // Same config for now
		// In real implementation, this would contain actual threshold keys
	}, nil
}

// CreateRealADKRFunc creates a function compatible with KeyRefreshFunc interface
func CreateRealADKRFunc(config RealADKRConfig, logger *log.Logger) KeyRefreshFunc {
	r := NewRealADKR(config, logger)
	return func(epoch int) (*adkr.KeyRefreshOutput, error) {
		return r.RunADKR(epoch)
	}
}

// CreateRealADKRFuncWithMetrics creates a function with metrics integration
func CreateRealADKRFuncWithMetrics(config RealADKRConfig, logger *log.Logger, metrics *Metrics) KeyRefreshFunc {
	r := NewRealADKR(config, logger)
	r.SetMetrics(metrics)
	return func(epoch int) (*adkr.KeyRefreshOutput, error) {
		return r.RunADKR(epoch)
	}
}
